from typing import (
    Any, Callable, ContextManager, Iterator, List, Literal, NamedTuple, Optional, Sequence, Tuple, TypeVar,
    Union, overload,
)

import torch
from torch import Generator, contiguous_format, inf, strided, SymInt
from torch.types import Device, Number, _bool, _complex, _device, _dtype, _float, _int, _layout, _qscheme, _size
from torchair._ge_concrete_graph import ge_apis as ge
from torchair._ge_concrete_graph.fx2ge_converter import declare_supported, register_fx_node_ge_converter
from torchair.ge._ge_graph import Tensor, TensorSpec, DataType
from torchair._ge_concrete_graph.supported_declaration import _TypedTensor, F32, F16, F64, I32, I16, I64, I8, U8, \
    BOOL, Support


@register_fx_node_ge_converter(torch.ops.npu.npu_moe_distribute_combine_setup.default)
def convert_npu_moe_distribute_combine_setup(
    expand_x: Tensor, 
    ep_send_counts: Tensor, 
    group_ep: str, 
    ep_world_size: int, 
    ep_rank_id: int, 
    moe_expert_num: int, 
    *, 
    expert_shard_type: int = 0, 
    shared_expert_num: int = 1, 
    shared_expert_rank_num: int = 0, 
    global_bs: int = 0, 
    comm_quant_mode: int = 0, 
    comm_type: int = 0,
    meta_outputs: TensorSpec = None
):

    return ge.MoeDistributeCombineSetup(expand_x=expand_x,
                                        ep_send_counts=ep_send_counts, 
                                        group_ep=group_ep,
                                        ep_world_size=ep_world_size,
                                        ep_rank_id=ep_rank_id,
                                        moe_expert_num=moe_expert_num, 
                                        expert_shard_type=expert_shard_type,
                                        shared_expert_num=shared_expert_num, 
                                        shared_expert_rank_num=shared_expert_rank_num, 
                                        global_bs=global_bs, 
                                        comm_quant_mode=comm_quant_mode, 
                                        comm_type=comm_type)